# 标定主程序
from import_data import *
from math import *


# 投影函数
def pcd_proj(points):
    velo = np.insert(points, 3, 1, axis=1).T
    cam = trans_mat * velo
    cam[:2] /= cam[2, :]
    u, v, tmp1 = cam
    return u, v


# 计算旋转矩阵
def rot_mat(sita_x, sita_y, sita_z):
    mat = np.matrix([[(cos(sita_z) * cos(sita_y)),
                      (cos(sita_z) * sin(sita_y) * sin(sita_x) - sin(sita_z) * cos(sita_x)),
                      (cos(sita_z) * sin(sita_y) * cos(sita_x) + sin(sita_z) * sin(sita_x))],
                     [(sin(sita_z) * cos(sita_y)),
                      (sin(sita_z) * sin(sita_y) * sin(sita_x) + cos(sita_z) * cos(sita_x)),
                      (sin(sita_z) * sin(sita_y) * cos(sita_x) - cos(sita_z) * sin(sita_x))],
                     [-sin(sita_y), cos(sita_y) * sin(sita_x), cos(sita_y) * cos(sita_x)]])
    return mat


# 相机内参
P2 = np.mat([[1916.240369950428,    0,                  0],
            [0,                     1913.278508450042,  0],
            [1021.529553590090,     532.6943887717723,  1],
            [0,                     0,                  0]]).T
# 图像尺寸
img_size = [1920, 1080]
# 最小像素误差平方和
min_err = 1e9
# 最优参数
min_11 = []
# 最小像素误差
err_list = []
# 区间划分精度
div_step = 6
# 平移参数初始值
x_head = -2
x_tail = 2
y_head = -2
y_tail = 2
z_head = 0
z_tail = 4
# 旋转标定参数初始值
xs_head = -0.1
xs_tail = 0.1
ys_head = -0.1
ys_tail = 0.1
zs_head = -0.1
zs_tail = 0.1
# 迭代次数
n_iter = 10
for i_iter in range(n_iter):
    print('----------iter_time: ' + str(i_iter + 1) + '----------')
    # 当前搜索区间
    print('x_range: ', end='')
    print(x_head, x_tail)
    print('y_range: ', end='')
    print(y_head, y_tail)
    print('z_range: ', end='')
    print(z_head, z_tail)
    print('x_sita_range: ', end='')
    print(xs_head, xs_tail)
    print('y_sita_range: ', end='')
    print(ys_head, ys_tail)
    print('z_sita_range: ', end='')
    print(zs_head, zs_tail)
    # 当前搜索区间点
    x_range = [x_head + (x_tail - x_head) / div_step * _i for _i in range(div_step)]
    y_range = [y_head + (y_tail - y_head) / div_step * _i for _i in range(div_step)]
    z_range = [z_head + (z_tail - z_head) / div_step * _i for _i in range(div_step)]
    sx_range = [xs_head + (xs_tail - xs_head) / div_step * _i for _i in range(div_step)]
    sy_range = [ys_head + (ys_tail - ys_head) / div_step * _i for _i in range(div_step)]
    sz_range = [zs_head + (zs_tail - zs_head) / div_step * _i for _i in range(div_step)]
    # 搜索主循环
    for sita_x0 in sx_range:
        for sita_y0 in sy_range:
            for sita_z0 in sz_range:
                # 旋转矩阵
                mat_rot = rot_mat(sita_x0, sita_y0, sita_z0)
                for x11 in x_range:
                    for y11 in y_range:
                        for z11 in z_range:
                            # 当前参数误差平方和
                            tmp_err = 0
                            # 当前平移矩阵
                            R0_rect = np.mat([[0, 0, 1, 0],
                                              [-1, 0, 0, 0],
                                              [0, -1, 0, 0],
                                              [x11, y11, z11, 1]]).T
                            # 当前投影矩阵
                            trans_mat = mat_rot * P2 * R0_rect
                            err_list_tmp = []
                            for i_data, pcd_data in enumerate(pcd_data_list):
                                img_data = img_data_list[i_data]
                                for data_label in pcd_data.keys():
                                    if data_label in img_data.keys():
                                        pcd_pos = np.mat(pcd_data[data_label])
                                        img_pos = [img_data[data_label][0] * img_size[0],
                                                   img_data[data_label][1] * img_size[1]]
                                        u0, v0 = pcd_proj(pcd_pos)
                                        w11, h11 = u0[0] - img_pos[0], v0[0] - img_pos[1]
                                        tmp_err += w11 ** 2 + h11 ** 2
                                        err_list_tmp.append([w11[0, 0], h11[0, 0]])
                            if tmp_err < min_err:
                                min_err = tmp_err
                                min_11 = [x11, y11, z11, sita_x0, sita_y0, sita_z0]
                                err_list = np.mat(err_list_tmp) * np.mat(np.eye(2))
    # 缩小搜索区间
    len_x = (x_tail - x_head) / 4
    x_head = max(min_11[0] - len_x, x_head)
    x_tail = min(min_11[0] + len_x, x_tail)
    len_y = (y_tail - y_head) / 4
    y_head = max(min_11[1] - len_y, y_head)
    y_tail = min(min_11[1] + len_y, y_tail)
    len_z = (z_tail - z_head) / 4
    z_head = max(min_11[2] - len_z, z_head)
    z_tail = min(min_11[2] + len_z, z_tail)
    len_xs = (xs_tail - xs_head) / 4
    xs_head = max(min_11[3] - len_xs, xs_head)
    xs_tail = min(min_11[3] + len_xs, xs_tail)
    len_ys = (ys_tail - ys_head) / 4
    ys_head = max(min_11[4] - len_ys, ys_head)
    ys_tail = min(min_11[4] + len_ys, ys_tail)
    len_zs = (zs_tail - zs_head) / 4
    zs_head = max(min_11[5] - len_zs, zs_head)
    zs_tail = min(min_11[5] + len_zs, zs_tail)
    print('opt_param: ', end='')
    print(min_11)
    print('min_sum_err: ', end='')
    print(min_err)
    
print('*********************')
print('opt_param: ', end='')
print(min_11)
print('min_sum_err: ', end='')
print(min_err)
print('trans_mat: ')
R0_rect = np.mat([[0, 0, 1, 0],
                  [-1, 0, 0, 0],
                  [0, -1, 0, 0],
                  [min_11[0], min_11[1], min_11[2], 1]]).T
mat_rot = rot_mat(min_11[3],  min_11[4], min_11[5])
R0_rect[:3, :3] = R0_rect[:3, :3] * mat_rot
print(R0_rect)

np.savetxt('err.txt', err_list, encoding='utf-8')
np.savetxt('calib.txt', R0_rect, encoding='utf-8')
